# from distutils.log import info
# from hashlib import new
# from re import sub
# import rdkit
import os
import numpy as np
import glob
import argparse
import subprocess
# import time

parser = argparse.ArgumentParser(description="path to lig structure file (mol/sdf) and the extension")

parser.add_argument('-lig_dir_path', type=str,
                    help="name of the ligand mol file (i.e. ligands)")
parser.add_argument('-lig_extension', type=str,
                    help="extension (i.e. .mol) of the ligand mol file")
parser.add_argument('-protein_pdb', type=str,
                    help="name of the protein pdb file with extension")

args = parser.parse_args()
lig_dir=args.lig_dir_path
lig_f_ext=args.lig_extension
ligfile_list=glob.glob(lig_dir+'/*'+lig_f_ext)
min1_mdp='min01.mdp'
min2_mdp='min02.mdp'
min3_mdp='min03.mdp'
nve_mdp='nve.mdp'
nvt_mdp='nvt.mdp'
npt_mdp='npt.mdp'
pro_struc='protein.pdb'

#ligfile_list=['ligands/1r.sdf'] # test case

water_shell=10 # 10 A thick water layer
box_type='dodecahedron'
ion_conc=0.15

water_box_data = {'spc216.gro': ['spc.itp', 'tip3p.itp', 'scpe.itp', 'tips3p.itp'], 'tip4p.gro': ['tip4p.itp'],
                  'tip5p.grp': ['tip5p.itp', 'tip5pe.itp']}

def print_logs(info_str):
    print('='*18)
    print(f'{info_str}')
    print('='*18)

for ligfile in ligfile_list:
    lig=ligfile[(len(lig_dir)+1):-len(lig_f_ext)]
    print(lig)
    lig_md_dir=lig+'_md/'
    if os.path.isdir(f'{lig_md_dir}'):
        subprocess.run(["rm","-rf",lig_md_dir])
    else:
        pass    
    subprocess.run(["mkdir",lig_md_dir])
    os.system(f'cp mdp/*.mdp {lig_md_dir}')
    os.system(f'cp ligands/{lig}/{lig}.mol {lig_md_dir}')
    os.system(f'cp ligands/{lig}/MOL.acpype/MOL_GMX.itp {lig_md_dir}{lig}.itp')
    os.system(f'cp ligands/{lig}/MOL.acpype/MOL_NEW.pdb {lig_md_dir}{lig}.pdb')
    
    print_logs(f'start seting up {lig_md_dir} sys vac')
    os.system(f'cp {pro_struc} {lig_md_dir}')
    os.system(f'printf "6\n1\n" | gmx pdb2gmx -f {lig_md_dir}{pro_struc} -p {lig_md_dir}protein.top -o {lig_md_dir}protein_step1.pdb -ignh -quiet')
    with open(f'{lig_md_dir}protein_step1.pdb','r') as file:
        contents=file.readlines()
    protein_pdb=''.join([line for line in contents 
                        if (line.find('ATOM')!=-1 or line.find('TER')!=-1)])
    with open(f'{lig_md_dir}protein_step2.pdb','w') as file:
        file.write(protein_pdb)
    with open(f'{lig_md_dir}{lig}.pdb') as file:
        contents=file.readlines()
    ligand_pdb=''.join([line for line in contents 
                        if (line.find('HETATM')!=-1 or (line.find('ATOM')!=-1 or line.find('END')!=-1))])
    with open(f'{lig_md_dir}complex.pdb', 'w') as file:
        file.write(protein_pdb)
        file.write('\n'+ligand_pdb)
    os.system(f'gmx editconf -f {lig_md_dir}complex.pdb -o {lig_md_dir}sys_vac.pdb -quiet')
    print_logs(f'finish setting up {lig_md_dir} sys vac')

    print_logs(f'start setting up sys vac topology')
    with open(f'{lig_md_dir}protein.top') as file:
        contents=file.readlines()
    top_info_list=[line for line in contents]
    with open(f'{lig_md_dir}{lig}.itp') as file:
        contents=file.readlines()
    lig_top_info_list=[line for line in contents]
    for (line_idx,line) in enumerate(lig_top_info_list):
        if '[ atomtypes ]' in line:
            lig_atp_start_idx=line_idx
        if '[ moleculetype ]' in line:
            lig_atp_stop_idx=line_idx

    with open(f'{lig_md_dir}lig.atp','w') as file:
        file.write(''.join(lig_top_info_list[lig_atp_start_idx:lig_atp_stop_idx]))

    with open(f'{lig_md_dir}lig.itp','w') as file:
        file.write(''.join(lig_top_info_list[lig_atp_stop_idx:]))

    # check if protein has multiple chains
    is_pro_top_written=True
    for line in top_info_list:
        if 'Include chain topologies' in line:
            is_pro_top_written=False
    # construct the sys_vac.top file in two different ways depending on weather protein has multiple chains
    if is_pro_top_written:
        for (line_idx,line) in enumerate(top_info_list):
            if '[ system ]' in line:
                top_info_list[line_idx-1]=top_info_list[line_idx-1]+f'\n; Include ligand topology\n#include "lig.itp"\n\n'
            elif '[ moleculetype ]' in line:
                top_info_list[line_idx-1]=top_info_list[line_idx-1]+f'\n; Include ligand atomtypes\n#include "lig.atp"\n\n'
            elif line_idx==len(top_info_list)-1:
                top_info_list[line_idx]=top_info_list[line_idx]+'MOL'+' '*12+'1\n'
            else:
                pass
    else:
        with open(f'{lig_md_dir}protein_Protein_chain_A.itp','r') as file:
            contents=file.readlines()
        pro_itp_info_list=[line for line in contents]
        for (line_idx,line) in enumerate(pro_itp_info_list):
            if '[ moleculetype ]' in line:
                pro_itp_info_list[line_idx-1]=pro_itp_info_list[line_idx-1]+f'\n; Include ligand atomtypes\n#include "lig.atp"\n\n'
        with open(f'{lig_md_dir}protein_Protein_chain_A_modify.itp','w') as file:
            file.write(''.join(pro_itp_info_list))
        for (line_idx,line) in enumerate(top_info_list):
            if '[ system ]' in line:
                top_info_list[line_idx-1]=top_info_list[line_idx-1]+f'\n; Include ligand topology\n#include "lig.itp"\n\n'
            elif '#include "protein_Protein_chain_A.itp"' in line:
                top_info_list[line_idx]=f'#include "protein_Protein_chain_A_modify.itp"\n'
            elif line_idx==len(top_info_list)-1:
                top_info_list[line_idx]=top_info_list[line_idx]+'MOL'+' '*12+'1\n'
            else:
                pass
    
    with open(f'{lig_md_dir}sys_vac.top','w') as file:
        file.write(''.join(top_info_list))
    print_logs(f"finish setting up sys vac topology")

    print_logs(f"solvate sys vac")
    box_size=water_shell/10.0
    os.system(f'gmx editconf -f {lig_md_dir}sys_vac.pdb -d {str(box_size)} -o {lig_md_dir}sys_vac.gro -bt {box_type}')
    os.system(f'cp {lig_md_dir}sys_vac.top {lig_md_dir}system.top')
    water_top=top_info_list[top_info_list.index('; Include water topology\n') + 1]
    solvent_box='spc216.gro' ### default tip3p water box gro file
    for water_box, water_models in water_box_data.items(): # check the top file to make sure
        for each_model in water_models:
            if water_top.find(each_model) != -1:
                solvent_box=water_box
    os.system(f'gmx solvate -cp {lig_md_dir}sys_vac.gro -cs {solvent_box} -o {lig_md_dir}sys_hydro.gro -p {lig_md_dir}system.top')
    print_logs(f"finish solvating sys vac")

    print_logs(f"add ions")
    with open(f'{lig_md_dir}genion_mdp.mdp','wb') as file:
        file.write(b'\n')
    genion_grompp_list=['gmx','grompp','-f',f'{lig_md_dir}genion_mdp.mdp',
                        '-c',f'{lig_md_dir}sys_hydro.gro','-p',f'{lig_md_dir}system.top',
                        '-o',f'{lig_md_dir}genion.tpr','-maxwarn','1',
                        '-po','genion_mdout.mdp']
    subprocess.run(genion_grompp_list)

    genion_list = ['gmx ','genion ', '-s ', f'{lig_md_dir}genion.tpr ',
                   '-p ',f'{lig_md_dir}system.top ','-o ',f'{lig_md_dir}system.pdb ',
                   '-conc ', str(ion_conc), ' -neutral ',
                   '-nname ','CL ','-pname ','NA']
    genion_args=['printf ','"SOL\n"',' | ']
    genion_command_list=genion_args+genion_list
    genion_commands=''.join(genion_command_list)
    os.system(genion_commands)
    print_logs(f"finish adding ions")
    os.system(f'printf "r protein mol\nq\n" | gmx make_ndx -f {lig_md_dir}system.pdb -o {lig_md_dir}index.ndx')
    print_logs(f"generate index file")
    
    print_logs(f"finish system {lig_md_dir}")

# print('start MD simulations')
# start=time.time()
# os.system('bash equ_run.sh')
# end=time.time()
# md_time_s=start-end
# md_time_min=md_time_s/60
# print(f'finished MD, and they take {md_time_min} min')

# os.system(f'bash run_gbsa_amber {lig_dir}')

##### run entropy estimation #####
# os.system(f'python get_entropy.py')
# end_final=time.time()
# total_t=end_final-start
# total_t_min=total_t/60
# print(f"Total time takes to run GBSA+entropy is {total_t_min} min")
